import os
import torch
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))
from torch import nn
from torchvision.models import resnet50, vit_b_16
from tqdm import tqdm
import numpy as np
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix
import torch.nn.functional as F
# from model_three.dataset_one_slice import MRISliceDataset

from model_effcientNet.dataset_SUVmax import MRIDataset
os.environ["CUDA_VISIBLE_DEVICES"]="0"
import torchvision.transforms as transforms
import pandas as pd
  # 非交互式后

class Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        base_model = resnet50(pretrained=False)
        encoder_layers = list(base_model.children())
        self.backbone = nn.Sequential(*encoder_layers[:9])
        
    def forward(self, x):
        return self.backbone(x)

class ModalitySpecificNet(nn.Module):
    def __init__(self, in_channels=5, num_classes=4):
        super().__init__()
        # 每个模态独立的主干网络
        self.backbone = Backbone() # 期待B，3，H，W的输入
        self.backbone.load_state_dict(torch.load("./pretrained_model/RadImageNet_pytorch/ResNet50.pt"))
        # 参数初始化
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, imgs, rois, labels):
        # 模态特定的处理
        feature = self.backbone(imgs).squeeze(-1).squeeze(-1)
        return feature

def train_model(model, train_loader, val_loader, criterion, optimizer, device, num_epochs=30):
    min_valid_loss = np.inf
    best_acc = 0

    def process_epoch(data_loader, is_train, epoch):
        model.train() if is_train else model.eval()
        stage = "训练" if is_train else "验证"
        
        total_loss = 0.0
        total_correct = 0
        total_samples = 0
        all_preds = []
        all_labels = []
        
        progress_bar = tqdm(data_loader, desc=f"{stage} Epoch {epoch+1}", leave=False)
        with torch.set_grad_enabled(is_train):
            for batch_idx, imgs, rois, labels, pid in enumerate(progress_bar):
                # 统一处理数据加载
                imgs = imgs.to(device)
                rois = rois.to(device)
                labels = labels.to(device)
                pid = pid.to(device)

                # 梯度清零仅在训练时进行
                if is_train:
                    optimizer.zero_grad()
                
                # 前向传播
                outputs = model(imgs, rois, labels, pid)
                loss = criterion(outputs, labels)
            
                # 反向传播和优化仅在训练时进行
                if is_train:
                    loss.backward()
                    optimizer.step()
                    
                total_loss += loss.item()
                # 统一计算指标
                probs = F.softmax(outputs, dim=1)
                preds = probs.argmax(dim=1)
                total_correct += (preds == labels).sum().item()
                total_samples += labels.size(0)
                
                # 收集预测结果和标签
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
                
                # 统一日志输出
                log_interval = 20 if is_train else 10
                if batch_idx % log_interval == 0:
                    progress_bar.write(
                        f"Batch {batch_idx}: 预测 {preds} | 真实 {labels.cpu().numpy()}"
                    )
        
        # 统一计算epoch指标
        epoch_loss = total_loss / total_samples
        epoch_acc = total_correct / total_samples
        
        # 验证阶段额外计算混淆矩阵
        if not is_train:
            cm = confusion_matrix(all_labels, all_preds, labels=[0, 1, 2, 3])
            cm_str = np.array2string(cm, separator='\t', 
                                    formatter={'int': lambda x: f"{x:3d}"},
                                    prefix='\t')
            print(f"验证混淆矩阵:\n{cm_str}\n")
        
        return epoch_loss, epoch_acc
    
    # 主训练循环
    for e in range(num_epochs):
        # 训练阶段
        train_loss, train_acc = process_epoch(train_loader, True, e)
        
        # 验证阶段
        val_loss, val_acc = process_epoch(val_loader, False, e)
        
        # 统一日志输出
        print(f'Epoch {e+1}')
        print(f'训练 Loss: {train_loss:.8f} Acc: {train_acc:.4f}')
        print(f'验证 Loss: {val_loss:.8f} Acc: {val_acc:.4f}')
        
        # 模型保存逻辑
        if val_acc > best_acc:
            print(f'验证损失下降 ({min_valid_loss:.6f} → {val_loss:.6f})')
            best_acc = val_acc
            torch.save(model.state_dict(), f'crc_fold_{val_acc:.4f}_best_model.pth')
            min_valid_loss = val_loss
        
        if (e+1) % 10 == 0:
            torch.save(model.state_dict(), f'crc_fold_{e+1}_model.pth')


def get_feature(model, train_loader, device):
    # 初始化数据结构
    df_list = []
    total_batches = len(train_loader)  # 获取总批次数量
    
    with torch.no_grad():
        model.eval()  # 设置为评估模式
        
        # 添加tqdm进度条
        for batch in tqdm(train_loader, total=total_batches, desc='Extracting Features', unit='batch'):
            imgs, rois, labels, pid = batch
            imgs, rois = imgs.to(device), rois.to(device)
            
            # 前向传播获取特征
            outputs = model(imgs, rois, labels)
            
            # 转换为numpy并处理
            features = outputs.cpu().numpy()
            pids = np.array(pid)
            labels = labels.cpu().numpy()
            
            # 构建批次DataFrame
            batch_df = pd.DataFrame({
                'PID': pids,
                'Label': labels
            })
            # 添加特征列
            feature_cols = [f'Feature_{i}' for i in range(features.shape[1])]
            batch_df = pd.concat([
                batch_df,
                pd.DataFrame(features, columns=feature_cols)
            ], axis=1)
            
            df_list.append(batch_df)

    # 合并所有数据
    final_df = pd.concat(df_list, ignore_index=True)
    
    # 保存到Excel（按PID排序）
    final_df.sort_values('PID').to_excel('features_by_pid.xlsx', index=False)
    print(f"特征已保存到 features_by_pid.xlsx，共 {len(final_df)} 条数据")

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = ModalitySpecificNet(in_channels=3, num_classes=4).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    data_transform = transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Resize((64, 64)),
                                    transforms.Normalize([0.5], [0.5]),
                                    transforms.Lambda(lambda x: x.squeeze(0)),
                                    ])
    train_dataset = MRIDataset(root_dir='D:/PyChrom/PythonProject/2.medical_image/DeepTrip/datasets/pT/three_label_data/train',
                            transform=data_transform)
                            
    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
    
    train_model(model, train_loader, criterion, optimizer, device, num_epochs=30)
    # get_feature(model, train_loader, device="cuda")
